#include "main/utils.h"
#include "lib/binary_tree.h"
using namespace std;

using TreeNode = BinaryTreeNode;

class Solution {
  public:
    void flatten(TreeNode* root) {
      if (!root) {
        return;
      }
      subflatten(root);
    }
    TreeNode* subflatten(TreeNode* root) {
      TreeNode* left = root->left;
      TreeNode* right = root->right;
      TreeNode* last = root;
      root->left = nullptr;
      if (left) {
        root->right = left;
        last = subflatten(left);
      }
      if (right) {
        last->right = right;
        last = subflatten(right);
      }
      return last;
    }

    void flatten1(TreeNode* root) {
      vector<TreeNode*> nodes;
      preorder(root, nodes);
      for (int i = 1; i < nodes.size(); ++i) {
        auto pre = nodes[i - 1];
        auto cur = nodes[i];
        pre->left = nullptr;
        pre->right = cur;
      }
    }
    void preorder(TreeNode* root, vector<TreeNode*>& nodes) {
      if (root) {
        nodes.push_back(root);
        preorder(root->left, nodes);
        preorder(root->right, nodes);
      }
    }

    void flatten2(TreeNode* root) {
      vector<TreeNode*> nodes_vec;
      stack<TreeNode*> nodes_stk;
      auto cur_node = root;
      while (cur_node || !nodes_stk.empty()) {
        while (cur_node) {
          nodes_vec.push_back(cur_node);
          nodes_stk.push(cur_node);
          cur_node = cur_node->left;
        }
        cur_node = nodes_stk.top();
        nodes_stk.pop();
        cur_node = cur_node->right;
      }
      for (int i = 1; i < nodes_vec.size(); ++i) {
        auto pre = nodes_vec[i - 1];
        auto cur = nodes_vec[i];
        pre->left = nullptr;
        pre->right = cur;
      }
    }

    void flatten3(TreeNode* root) {
      if (!root) {
        return;
      }
      stack<TreeNode*> stk;
      stk.push(root);
      TreeNode* pre = nullptr;
      while (!stk.empty()) {
        auto cur = stk.top();
        stk.pop();
        if (pre) {
          pre->left = nullptr;
          pre->right = cur;
        }
        if (cur->right) {
          stk.push(cur->right);
        }
        if (cur->left) {
          stk.push(cur->left);
        }
        pre = cur;
      }
    }

    void flatten4(TreeNode* root) {
      TreeNode* cur_node = root;
      while (cur_node) {
        auto left = cur_node->left;
        if (left) {
          auto next = left;
          auto right_pre = next;
          while (right_pre->right) {
            right_pre = right_pre->right;
          }
          right_pre->right = cur_node->right;
          cur_node->right = left;
          cur_node->left = nullptr;
        }
        cur_node = cur_node->right;
      }
    }
};

int main() {
  vector<int> vec = {1, 2, 5, 3, 4, 6};
  Solution su;

  auto root = CreateBinaryTree(vec);
  su.flatten(root);
  PrintBinaryTree(root);

  auto root1 = CreateBinaryTree(vec);
  su.flatten1(root1);
  PrintBinaryTree(root1);

  auto root2 = CreateBinaryTree(vec);
  su.flatten2(root2);
  PrintBinaryTree(root2);

  auto root3 = CreateBinaryTree(vec);
  su.flatten3(root3);
  PrintBinaryTree(root3);

  auto root4 = CreateBinaryTree(vec);
  su.flatten4(root4);
  PrintBinaryTree(root4);
}
